﻿using System;
using System.Collections.Generic;
using System.Text;
using EnvDTE;
using EnvDTE80;
using System.Collections;
using Microsoft.VisualStudio.VCProjectEngine;
using System.IO;
using System.Net.Sockets;
using System.Threading;
using Microsoft.VisualStudio.Shell;

namespace Gauss.Isp
{
    public class Debugger
    {
        private List<Envelope> _transitions;
        private int _currentTransitionIndex;
        private int _numberProcesses;
        private DTE2 _dte;
        private _dispDebuggerEvents_OnEnterDesignModeEventHandler _debuggerExitHandler;
        private _dispDebuggerEvents_OnEnterRunModeEventHandler _debuggerRunHandler;
        private _dispDebuggerEvents_OnEnterBreakModeEventHandler _debuggerBreakHandler;
        private bool _propertiesSaved;
        private Dictionary<string, object> _propertiesStore;
        private DebuggerEvents _debuggerEvents;
        private Project _startupProject;
        private ProjectHandler _startupProjectHandler;
        private VCDebugSettings _debugSettings;
        private object _savedStartupProjects;
        private bool _debuggerRunning;
        private List<Project> _projects;
        private ToolsOptions _toolsOptions;
        private DebuggerToolWindow _toolWindow;
        private bool _toolWindowInterleavingSet;
        private int _currentLockedRank;

        public Debugger(IEnumerable<Envelope> transitions, Envelope transition, int numProcs, List<Project> projects, DTE2 dte, ToolsOptions toolsOptions, DebuggerToolWindow toolWindow)
        {
            _transitions = new List<Envelope>(transitions);
            _transitions.Sort(
                delegate(Envelope a, Envelope b)
                {
                    if (a.OrderIndex == b.OrderIndex)
                    {
                        if (a.Rank != b.Rank)
                            return a.Rank.CompareTo(b.Rank);
                        else
                            return a.Index.CompareTo(b.Index);
                    }
                    else
                        return a.OrderIndex.CompareTo(b.OrderIndex);
                });
            _currentTransitionIndex = _transitions.IndexOf(transition);
            _numberProcesses = numProcs;
            _dte = dte;
            _propertiesStore = new Dictionary<string, object>();
            _debuggerEvents = _dte.Events.DebuggerEvents;

            // See what project we just start.
            _projects = projects;
            if (projects == null || projects.Count == 0)
            {
                throw new ApplicationException("No projects are specified as the startup project.");
            }
            _startupProject = projects[0];
            _startupProjectHandler = new ProjectHandler(_startupProject);
            _debugSettings = _startupProjectHandler.GetDebugSettings();
            _toolsOptions = toolsOptions;
            _toolWindow = toolWindow;

            // Set up the event handlers for debugger Run, Break, and Exit.
            _debuggerRunning = _dte.Debugger.CurrentMode != dbgDebugMode.dbgDesignMode;
            _debuggerExitHandler = new _dispDebuggerEvents_OnEnterDesignModeEventHandler(Debugger_Exit);
            _debuggerRunHandler = new _dispDebuggerEvents_OnEnterRunModeEventHandler(Debugger_Run);
            _debuggerBreakHandler = new _dispDebuggerEvents_OnEnterBreakModeEventHandler(Debugger_Break);
            _debuggerEvents.OnEnterDesignMode += _debuggerExitHandler;
            _debuggerEvents.OnEnterRunMode += _debuggerRunHandler;
            _debuggerEvents.OnEnterBreakMode += _debuggerBreakHandler;

            // Interface with the tool window to control the stepping.
            _toolWindow.NextTransitionClick += new EventHandler(NextTransition_Click);
            _toolWindow.LockedRankChanged += new EventHandler<LockedRankChangedEventArgs>(LockedRank_Changed);
            _toolWindow.TransitionClicked += new EventHandler<TransitionClickedEventArgs>(Transition_Clicked);
        }

        private void Transition_Clicked(object sender, TransitionClickedEventArgs e)
        {
            int newIndex = _transitions.IndexOf(e.Transition);
            if (newIndex > _currentTransitionIndex)
            {
                _currentTransitionIndex = newIndex;
                lock (_toolWindow)
                {
                    Monitor.Pulse(_toolWindow);
                }
            }
        }

        private void NextTransition_Click(object sender, EventArgs e)
        {
            lock (_toolWindow)
            {
                Monitor.Pulse(_toolWindow);
            }
        }

        private void LockedRank_Changed(object sender, LockedRankChangedEventArgs e)
        {
            _currentLockedRank = e.LockedRank;
        }

        public void Start()
        {
            // Already running, stop.
            if (_debuggerRunning)
            {
                return;
            }

            string mpishimBatchPath = null;
            DebuggerServer debuggerServer = null;

            // Save the original debugger settings.
            _propertiesStore.Clear();
            _propertiesStore["DebuggerFlavor"] = _debugSettings.DebuggerFlavor;
            _propertiesStore["MPIRunCommand"] = _debugSettings.MPIRunCommand;
            _propertiesStore["MPIRunArguments"] = _debugSettings.MPIRunArguments;
            _propertiesStore["MPIRunWorkingDirectory"] = _debugSettings.MPIRunWorkingDirectory;
            _propertiesStore["ApplicationCommand"] = _debugSettings.ApplicationCommand;
            _propertiesStore["ApplicationArguments"] = _debugSettings.ApplicationArguments;
//            _propertiesStore["ShimCommand"] = _debugSettings.ShimCommand;
            _propertiesStore["MPIAcceptMode"] = _debugSettings.MPIAcceptMode;
            _propertiesStore["MPIAcceptFilter"] = _debugSettings.MPIAcceptFilter;
            _propertiesStore["DebuggerType"] = _debugSettings.DebuggerType;
            _propertiesSaved = true;

            try
            {
                // Get the path to the ISP debugging helper.
                FileInfo ispDebugCodePath = _toolsOptions.Get_IspDebugCodePath;
                if (ispDebugCodePath == null)
                {
                    throw new FileNotFoundException(Utils.StringFormat(Resources.Message_CannotFindFile, "isp_debug.c"));
                }

                // Get the projects needed to be compiled.
                List<ProjectHandler> projects = new List<ProjectHandler>(_projects.Count - 1);
                for (int i = 1; i < _projects.Count; i++)
                {
                    projects.Add(new ProjectHandler(_projects[i]));
                }

                // Rebuild the projects.
                _dte.ExecuteCommand("File.SaveAll", "");

                // Set all of the compiler and linker paths correctly.
                foreach (ProjectHandler ph in projects)
                {
                    ph.AddFileToProject(ispDebugCodePath.FullName);
                    ph.AddLinkerDependency("Ws2_32.lib");
                    ph.Save();
                }

                // Build all of the projects.
                try
                {
                    foreach (ProjectHandler ph in projects)
                    {
                        IspBuildThread buildThread = new IspBuildThread(_dte, ph.ActiveConfigName,
                                    ph.UniqueName);
                        buildThread.Build();
                    }
                }
                catch
                {
                    // If failed, restore the project settings.
                    foreach (ProjectHandler ph in projects)
                    {
                        ph.RestoreProjectFiles();
                        ph.RestoreLinkerDependencies();
                        ph.Save();
                    }
                    throw;
                }

                // Find the executable that is being generated.
                string executable = _startupProjectHandler.GetExecutable();
                if (string.IsNullOrEmpty(executable))
                {
                    throw new FileNotFoundException(Resources.Message_CouldNotFindExecutable);
                }

                // Find the path to the working directory.
                string workingDirectoryPath = _startupProjectHandler.GetWorkingDirectory();
                workingDirectoryPath.TrimEnd(new char[] { '\\' });
                workingDirectoryPath = workingDirectoryPath.Substring(0, workingDirectoryPath.Length - 1); 

                // Set up a listening server on the first port available.
                int port = 9999;
                do
                {
                    try
                    {
                        debuggerServer = new DebuggerServer(port, _numberProcesses);
                    }
                    catch (SocketException)
                    {
                        debuggerServer = null;
                        port++;
                    }
                } while (debuggerServer == null);

                // This is the batch script executes mpiexec. Need to do it
                // in a script, because there are some strange Visual Studio
                // issues with spaces in the path.
                mpishimBatchPath = Path.GetTempPath() + Guid.NewGuid().ToString() + ".bat";
                using (StreamWriter file = new StreamWriter(mpishimBatchPath))
                {
                    file.WriteLine("@echo on");
                    file.WriteLine();
                    file.WriteLine("@rem cmdLine will store the arguments to this bat file, without the first two arguments.");
                    file.WriteLine("@rem The first two arguments are the mpishim path and the computer name. Those arguments");
                    file.WriteLine("@rem are already in the string below.");
                    file.WriteLine("set cmdLine=");
                    file.WriteLine("shift");
                    file.WriteLine("shift");
                    file.WriteLine(":VarLoop");
                    file.WriteLine("  if (%1)==() goto VarLoopEnd");
                    file.WriteLine("  set cmdLine=%cmdLine%%1 ");
                    file.WriteLine("  shift");
                    file.WriteLine("goto VarLoop");
                    file.WriteLine(":VarLoopEnd");
                    file.WriteLine();
                    file.WriteLine(Utils.StringFormat("\"{0}\" -n {1} -wdir \"{2}\" -env ISP_HOSTNAME {3} -env ISP_PORT {4} \"{5}\" {3} %cmdLine%",
                                                "mpiexec.exe",
                                                _numberProcesses,
                                                workingDirectoryPath,
                                                "127.0.0.1",
                                                port.ToString(),
                                                _toolsOptions.Get_MpishimPath)); //.Replace("\\", "\\\\")));
                }

                // Set the new settings for debugging.
                _debugSettings.DebuggerFlavor = eDebuggerTypes.eClusterDebugger;
                _debugSettings.MPIRunCommand = mpishimBatchPath;
                _debugSettings.MPIRunArguments = "";
                _debugSettings.MPIRunWorkingDirectory = "";
                _debugSettings.ApplicationCommand = executable; //.Replace("\\", "\\\\");
                _debugSettings.ApplicationArguments = _startupProjectHandler.GetProgramArguments();
                _debugSettings.ShimCommand = @"C:\Program Files (x86)\Microsoft Visual Studio 10.0\Common7\IDE\Remote Debugger\x64\mpishim100.exe"; //TODO: is this correct?  It was originally "mpishim"
                _debugSettings.MPIAcceptMode = enumMPIAcceptModes.eMPIAcceptSubnets;
                _debugSettings.MPIAcceptFilter = "";
                _debugSettings.DebuggerType = TypeOfDebugger.DbgNativeOnly;

                _toolWindowInterleavingSet = false;

                // Save the startup project, switch to the selected project, and start the debugger.
                _savedStartupProjects = _dte.Solution.SolutionBuild.StartupProjects;
                if (((object[])_savedStartupProjects).Length == 1)
                    _savedStartupProjects = ((object[])_savedStartupProjects)[0];
                _dte.Solution.SolutionBuild.StartupProjects = _startupProject.UniqueName;
                _dte.Solution.SolutionBuild.Debug();

                // Wait for the clients to connect.
                debuggerServer.WaitForClients();

                // Set the rank to pid mapping for the debugger toolwindow.
                _toolWindow.SetRankToPid(debuggerServer.GetRankToPidMapping());

                // After the clients connect, we can set all of the changed
                // properties back to the original 
                ResetProperties();
                foreach (ProjectHandler ph in projects)
                {
                    ph.RestoreProjectFiles();
                    ph.RestoreLinkerDependencies();
                    ph.Save();
                }

                // Handles the stepping of the debugged processes according to
                // what ISP says.
                DebuggerStepper(debuggerServer);

                // Tell the clients to no longer communicate with the server.
                for (int i = 0; i < _numberProcesses; i++)
                {
                    debuggerServer.WriteString(i, "2");
                }
            }
            finally
            {
                // Close the connections.
                if (debuggerServer != null)
                {
                    debuggerServer.CloseConnections();
                }

                // After the debugger stepping is done, can delete the .bat file.
                if (mpishimBatchPath != null)
                {
                    try
                    {
                        File.Delete(mpishimBatchPath);
                    }
                    catch
                    {
                    }
                }

                ResetProperties();
            }
        }

        private void DebuggerStepper(DebuggerServer debuggerServer)
        {
            if (_currentTransitionIndex < 0)
                _currentTransitionIndex = 0;
            int previousRank = -1;

            // Step to the current transition.
            for (int i = 0; i < _transitions.Count; i++)
            {
                Envelope e = _transitions[i];

                // First, step to the current transition without interruption,
                // and then require the continue button to be pressed.
                if (i > _currentTransitionIndex &&
                    (_currentLockedRank < 0 || _currentLockedRank == previousRank))
                {
                    lock (_toolWindow)
                    {
                        Monitor.Wait(_toolWindow);
                    }
                    try
                    {
                        _dte.Debugger.Go(false);
                    }
                    catch { }
                }

                _toolWindow.SetTransition(e);
                string call = debuggerServer.PeakNextEnvelope(e.Rank);
                if (call == null)
                {
                    return;
                }
                if (call != e.MpiCallName)
                {
                    Utils.MessageBox(null, "Expected " + e.MpiCallName + ", got " + call + ".");
                }

                debuggerServer.Step(e.Rank, i == _transitions.Count - 1 ||
                                            (i >= _currentTransitionIndex &&
                                             (_currentLockedRank < 0 || _currentLockedRank == e.Rank)));
                if (e.SourceRankSet)
                //if (e.MpiCallName == "MPI_Recv" || e.MpiCallName == "MPI_Irecv")
                {
                    if (e.MatchRank < 0)
                        debuggerServer.WriteString(e.Rank, e.SourceRank.ToString());
                    else
                        debuggerServer.WriteString(e.Rank, e.MatchRank.ToString());
                }

                previousRank = e.Rank;
            }
        }

        private void Debugger_Exit(dbgEventReason Reason)
        {
            _debuggerRunning = false;

            try
            {
                Microsoft.VisualStudio.Shell.Interop.IVsWindowFrame windowFrame = (Microsoft.VisualStudio.Shell.Interop.IVsWindowFrame)_toolWindow.Frame;
                windowFrame.CloseFrame((uint)vsSaveChanges.vsSaveChangesNo);
            }
            catch
            {
            }
        }

        private void Debugger_Run(dbgEventReason Reason)
        {
            _debuggerRunning = true;

            try
            {
                if (!_toolWindowInterleavingSet)
                {
                    _toolWindowInterleavingSet = true;
                    _currentLockedRank = -1;
                    _toolWindow.SetInterleaving(_transitions, _numberProcesses);
                    Microsoft.VisualStudio.Shell.Interop.IVsWindowFrame windowFrame = (Microsoft.VisualStudio.Shell.Interop.IVsWindowFrame)_toolWindow.Frame;
                    windowFrame.Show();
                }
            }
            catch
            {
            }
        }

        private void Debugger_Break(dbgEventReason Reason, ref dbgExecutionAction ExecutionAction)
        {
            Debugger_Run(Reason);
        }

        private void ResetProperties()
        {
            lock (_propertiesStore)
            {
                if (!_propertiesSaved)
                {
                    return;
                }

                // Restore the startup project.
                _dte.Solution.SolutionBuild.StartupProjects = _savedStartupProjects;

                // Restore the debugger settings.
                _debugSettings.DebuggerFlavor = (eDebuggerTypes) _propertiesStore["DebuggerFlavor"];
                _debugSettings.MPIRunCommand = _propertiesStore["MPIRunCommand"] as string;
                _debugSettings.MPIRunArguments = _propertiesStore["MPIRunArguments"] as string;
                _debugSettings.MPIRunWorkingDirectory = _propertiesStore["MPIRunWorkingDirectory"] as string;
                _debugSettings.ApplicationCommand = _propertiesStore["ApplicationCommand"] as string;
                _debugSettings.ApplicationArguments = _propertiesStore["ApplicationArguments"] as string;
                _debugSettings.ShimCommand = _propertiesStore["ShimCommand"] as string;
                _debugSettings.MPIAcceptMode = (enumMPIAcceptModes)_propertiesStore["MPIAcceptMode"];
                _debugSettings.MPIAcceptFilter = _propertiesStore["MPIAcceptFilter"] as string;
                _debugSettings.DebuggerType = (TypeOfDebugger)_propertiesStore["DebuggerType"];

                _propertiesSaved = false;
            }
        }
    }

    public class DebuggerServer
    {
        private int _count;

        private string[] _nextEnvelopes;
        private bool[] _nextEnvelopesFilled;

        private TcpListener _listener;
        private TcpClient[] _clients;
        private StreamReader[] _streamReaders;
        private StreamWriter[] _streamWriters;

        private Dictionary<int, int> _rankToId;
        private Dictionary<int, int> _idToPid;
        private Dictionary<int, int> _rankToPid;

        public DebuggerServer(int port, int count)
        {
            // Start listener on the specified port.
            _listener = new TcpListener(System.Net.IPAddress.Any, port);
            _listener.Start(count);

            _count = count;
            _clients = null;
            _streamReaders = null;
            _streamWriters = null;
            _rankToId = null;
            _idToPid = null;
        }

        public bool WaitForClients()
        {
            if (_listener == null)
            {
                return false;
            }

            _clients = new TcpClient[_count];
            _streamReaders = new StreamReader[_count];
            _streamWriters = new StreamWriter[_count];
            _nextEnvelopes = new string[_count];
            _nextEnvelopesFilled = new bool[_count];
            _rankToId = new Dictionary<int, int>();
            _idToPid = new Dictionary<int, int>();
            _rankToPid = new Dictionary<int, int>();

            for (int i = 0; i < _count; i++)
            {
                // Accept the connection and create the stream readers and writers.
                _clients[i] = _listener.AcceptTcpClient();
                _streamReaders[i] = new StreamReader(_clients[i].GetStream());
                _streamWriters[i] = new StreamWriter(_clients[i].GetStream());
                _streamWriters[i].AutoFlush = true;
                _streamWriters[i].NewLine = "\n";
                _nextEnvelopes[i] = null;
                _nextEnvelopesFilled[i] = false;

                // Read the rank and pid.
                int rank = int.Parse(_streamReaders[i].ReadLine());
                int pid = int.Parse(_streamReaders[i].ReadLine());
                _rankToId.Add(rank, i);
                _idToPid.Add(i, pid);
                _rankToPid.Add(rank, pid);
            }
            _listener.Stop();
            _listener = null;
            return true;
        }

        public void CloseConnections()
        {
            if (_clients == null)
                return;

            for (int i = 0; i < _clients.Length; i++)
            {
                _clients[i].Close();
            }

            _clients = null;
            _rankToId = null;
            _idToPid = null;
        }

        public string PeakNextEnvelope(int rank)
        {
            int id = _rankToId[rank];
            if (!_nextEnvelopesFilled[id])
            {
                _nextEnvelopes[id] = ReadString(rank);
                _nextEnvelopesFilled[id] = true;
            }
            return _nextEnvelopes[id];
        }

        public void Step(int rank, bool breakProcess)
        {
            int id = _rankToId[rank];
            if (!_nextEnvelopesFilled[id])
                ReadString(rank);
            WriteString(rank, breakProcess ? "1" : "0");
            _nextEnvelopesFilled[id] = false;
        }

        public int RankToPid(int rank)
        {
            return _idToPid[_rankToId[rank]];
        }

        public Dictionary<int, int> GetRankToPidMapping()
        {
            return _rankToPid;
        }

        private string ReadString(int rank)
        {
            try
            {
                int id = _rankToId[rank];
                if (!_streamReaders[id].EndOfStream)
                    return _streamReaders[id].ReadLine();
                else
                    return null;
            }
            catch
            {
                return null;
            }
        }

        public void WriteString(int rank, string str)
        {
            try
            {
                int id = _rankToId[rank];
                _streamWriters[id].WriteLine(str);
            }
            catch
            {
            }
        }
    }
}
